# -*- coding: utf-8 -*-
import os
import random
from pathlib import Path
from glob import glob

import numpy as np
from PIL import Image

import torch
from torch import nn
from torchvision import transforms

import seaborn as sns
import matplotlib
matplotlib.use("Qt5Agg")  # 指定使用 PyQt5
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
from tsne_function import DSResNet_extract_features

# 方案 A：系统里有“SimHei”（黑体）
plt.rcParams['font.sans-serif'] = ['SimHei']         # 显示中文
plt.rcParams['axes.unicode_minus'] = False           # 正常显示负号

# （可选）让 seaborn 也用同样设置
sns.set_theme(font='SimHei', rc={'axes.unicode_minus': False})

# ======== 1) 基本配置 ========
# four_data 的根目录（包含 0/1/2/3 子目录）
ROOT = Path("four_data_1750")     # 如需绝对路径：Path(r"D:\pycharm_project\Bearing_inspection\Be_legend\data_handle\four_data")
# 模型权重路径
MODEL_PATH = Path(r"D:\pycharm_project\Bearing_inspection\Be_legend\data_handle\model_save\DN4Net_best.pth")

# 每个类别采样数量
K_PER_CLASS = 1000
# 随机种子（可复现）
random.seed(2033)

# ======== 2) 区间定义与展开 ========
intervals = [
    [[3341, 6183]],
    [[0, 478], [2033, 2867], [3224, 3340]],
    [[956, 2032], [2510, 2628]],
    [[479, 955], [2271, 2509], [2868, 3223]],
]

# 展开为每类完整索引列表（端点包含）
expanded_intervals = [
    [i for (start, end) in cls for i in range(start, end + 1)]
    for cls in intervals
]

# 从每个类别随机采样 K_PER_CLASS 个索引（不足则全取）
sampled_indices = [
    random.sample(values, min(K_PER_CLASS, len(values)))
    for values in expanded_intervals
]

# ======== 3) 预处理与文件匹配 ========
# 三种模态（用于核对/缓存；推理只用 GADF 与 EMTF）
MODALITIES = ["CWT", "EMTF", "GADF"]

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485], std=[0.229])
])

# 用于缓存已读取的张量，避免重复 IO
# key: "00012_类名_模态.png"  value: Tensor[1,1,128,128]
tensor_cache = {}

def load_tensor_if_exists(cls_id: int, modality: str, idx: int):
    """
    尝试从 four_data/{cls}/{modality}/{idx:05d}.png 读取单通道图像并转 Tensor。
    若不存在，返回 None。
    """
    filename = f"{idx:05d}.png"
    path = ROOT / f"{cls_id}" / modality / filename
    if not path.exists():
        return None, str(path)

    key = f"{idx:05d}_{cls_id}_{modality}.png"
    if key in tensor_cache:
        return tensor_cache[key], None

    img = Image.open(path).convert("L")
    tensor = transform(img).unsqueeze(0)  # [1,1,128,128]
    tensor_cache[key] = tensor
    return tensor, None

# ======== 4) 模型加载 ========
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
from Neural_Networks import P2D_GE_Net, CNNNet, DN4Net, DSResNet_GADF  # 确保类名与权重匹配

model = DN4Net(in_channels=1, num_classes=4).to(device)
# 严格匹配加载
state = torch.load(MODEL_PATH, map_location=device)
model.load_state_dict(state)
model.eval()

# ======== 5) 推理：逐类逐样本，从 GADF + EMTF 两路输入预测 ========
true_labels, pred_labels = [], []
missing_files = []

with torch.no_grad():
    for cls_id, idx_list in enumerate(sampled_indices):
        for idx in idx_list:
            # 只用 GADF + EMTF（与你模型 forward(gadf, emtf) 一致）
            t_gadf, miss_g = load_tensor_if_exists(cls_id, "GADF", idx)
            t_emtf, miss_e = load_tensor_if_exists(cls_id, "EMTF", idx)
            t_cwt, miss_c = load_tensor_if_exists(cls_id, "CWT", idx)

            # 若缺任一模态，则跳过并记录
            if t_gadf is None or t_emtf is None:
                if miss_g: missing_files.append(miss_g)
                if miss_e: missing_files.append(miss_e)
                continue

            # 推理
            logits = model(t_gadf.to(device))
            y_pred = torch.argmax(logits, dim=1).item()

            true_labels.append(cls_id)
            pred_labels.append(y_pred)

# ======== 6) 打印缺失文件（若有） ========
if missing_files:
    print(f"[WARN] 缺失文件 {len(missing_files)} 个，示例：")
    for p in missing_files[:10]:
        print("  -", p)

print("预测样本数：", len(pred_labels))

# ======== 7) 混淆矩阵与可视化 ========
if len(true_labels) > 0:
    cm = confusion_matrix(true_labels, pred_labels, labels=[0,1,2,3])
    # 行归一化
    with np.errstate(all='ignore'):
        cm_norm = cm.astype(float) / cm.sum(axis=1, keepdims=True)

    plt.figure(figsize=(8, 6), dpi=300)
    sns.heatmap(
        cm_norm, annot=True, fmt=".2f", cmap="Blues", cbar=True,
        xticklabels=["0", "1", "2", "3"],
        yticklabels=["0", "1", "2", "3"]
    )
    plt.xlabel("预测类别")
    plt.ylabel("真实类别")
    plt.title("DN4Net 混淆矩阵")
    plt.tight_layout()
    out_png = "DN4Net_confusion_matrix.png"
    plt.savefig(out_png, dpi=300)
    plt.show()
    print(f"[OK] 混淆矩阵已保存：{out_png}")
else:
    print("[ERROR] 未得到有效预测（可能所有样本都缺少 GADF/EMTF 任意一模态）。请检查 four_data 目录结构与文件是否齐全。")

# ======== 8) t-SNE 可视化 ========
from sklearn.manifold import TSNE
import seaborn as sns

# 1) 收集能用的样本（同时存在 GADF 与 EMTF）
gadf_list, emtf_list, y_list = [], [], []
for cls_id, idx_list in enumerate(sampled_indices):
    for idx in idx_list:
        t_gadf, _ = load_tensor_if_exists(cls_id, "GADF", idx)
        t_emtf, _ = load_tensor_if_exists(cls_id, "EMTF", idx)
        if t_gadf is None or t_emtf is None:
            continue
        gadf_list.append(t_gadf)   # [1,1,128,128]
        emtf_list.append(t_emtf)   # [1,1,128,128]
        y_list.append(cls_id)

if len(y_list) == 0:
    print("[ERROR] t-SNE：没有可用样本（缺少 GADF/EMTF）。")
else:
    # 2) 拼批次
    Xg = torch.cat(gadf_list, dim=0)   # [N,1,128,128]
    Xe = torch.cat(emtf_list, dim=0)   # [N,1,128,128]
    Y  = torch.tensor(y_list, dtype=torch.long)

    ds  = torch.utils.data.TensorDataset(Xg, Xe, Y)
    dl  = torch.utils.data.DataLoader(ds, batch_size=64, shuffle=False)

    # 3) 通过前向钩子从 fc1 抽取特征（16 维）
    feats = []
    def _hook(module, inp, out):
        feats.append(out.detach().cpu())

    handle = model.classifier.register_forward_hook(_hook)
    model.eval()

    with torch.no_grad():
        for xg, xe, _y in dl:
            xg = xg.to(device)
            # xe = xe.to(device)
            _ = model(xg)  # 触发钩子捕获 fc1 输出

    handle.remove()
    features = torch.cat(feats, dim=0).numpy()   # [N, 16]
    labels   = Y.numpy()

    # 4) t-SNE
    tsne = TSNE(n_components=2, perplexity=30, n_iter=1000, random_state=42)
    features_2d = tsne.fit_transform(features)

    # 5) 绘图（背景设为纯白）
    plt.figure(figsize=(9, 7), dpi=300)
    sns.set_style("white")  # ✅ 设置纯白背景
    plt.gca().set_facecolor("#ffffff")  # ✅ 明确设定背景为白色

    class_names = ['normal', 'inner_ring', 'outer_ring', 'rolling_element']
    colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728']  # 蓝、橙、绿、红，与上图一致

    for cls, color in zip(np.unique(labels), colors):
        idxs = labels == cls
        plt.scatter(features_2d[idxs, 0], features_2d[idxs, 1],
                    label=class_names[cls], s=12, alpha=0.8, color=color)

    plt.legend(title="Class", fontsize=10)
    plt.title("t-SNE of DN4Net fc1 Features", fontsize=13)
    plt.tight_layout()
    plt.savefig("DN4Net_tsne_white.png", dpi=300, facecolor="#ffffff")  # ✅ 保存为白背景
    plt.show()

    print("[OK] t-SNE 图已保存：DN4Net_tsne_white.png")
